import numpy as np
from math import sqrt
from collections import Counter
from .metrics import accuracy


class KnnClassifer:
    def __init__(self, k):
        """
        初始化数据集
        :param k: 近邻个数
        """
        self.k = k
        self._X_train = None
        self._y_train = None

    def fit(self, X_train, y_train):
        """
        注入训练数据集
        :param X_train:  data
        :param y_train:  target
        :return:
        """
        assert X_train is not None and y_train is not None
        assert self._X_train is None and self._y_train is None
        assert X_train.shape[0] == y_train.shape[0]

        self._X_train = X_train
        self._y_train = y_train
        return self

    def predict(self, X_predict):
        assert self._X_train is not None and self._y_train is not None and X_predict is not None
        assert self._X_train.shape[1] == X_predict.shape[1]

        y_predict = [self._predict(x) for x in X_predict]
        return np.array(y_predict)

    def _predict(self, x):
        distance = [sqrt(np.sum((x_train - x) ** 2)) for x_train in self._X_train]
        nearest = np.argsort(distance)
        topK_y = [self._y_train[i] for i in nearest[:self.k]]
        vote = Counter(topK_y)
        return vote.most_common()[0][0]

    def score(self, X_test, y_test):
        y_predict = self.predict(X_test)
        return accuracy(y_test, y_predict)

    def __repr__(self):
        return "KNN MODEL(k=%d)" % self.k
